import torch
import os
import numpy as np
from torchvision.utils import save_image
import cfg
from PIL import Image
import cv2




if __name__ == '__main__':
    root_path='/home/ubuntu/PycharmProjects/u2net/data/test0'
    save_path='/home/ubuntu/PycharmProjects/u2net/data/test0_crop'
    img_list=os.listdir(os.path.join(root_path,"mask"))
    for name in img_list:
        img_path=os.path.join(root_path,"images",name)
        mask_path=os.path.join(root_path,"mask",name)
        assert(os.path.exists(img_path))
        assert(os.path.exists(mask_path))

        img=cv2.imread(img_path)
        mask=cv2.imread(mask_path)

        h,w,=img.shape[0],img.shape[1]
        step = cfg.img_size
        rh, rw = h, w

        pad_h = step - h if h < step else 0
        pad_w = step - w if w < step else 0
        m =torch.nn.ConstantPad2d((0, pad_w, 0, pad_h),0.)
        imgs=m(torch.Tensor(img)).numpy().astype(np.uint8)
        masks=m(torch.Tensor(mask)).numpy().astype(np.uint8)

        h, w = imgs.shape[0],imgs.shape[1]
        if(0):Image.fromarray(imgs[...,::-1]).show()

        clip_num=0
        for i in range(0, h, step):
            for j in range(0, w, step):
                if h - i < step:
                    i = h - step
                if w - j < step:
                    j = w - step
                clip_img = imgs[ i:i + step, j:j + step]
                clip_mask = masks[ i:i + step, j:j + step]
                assert(np.max(clip_mask)<=255.),"ddddddddddddddddddd"
                if(clip_img.shape[0]!=cfg.img_size and clip_img.shape[1]!=cfg.img_size): continue
                


                if(np.sum(255-clip_mask)>255.*3*1000):
                    # cv2.imshow("mask",np.uint8(clip_mask))
                    # cv2.waitKey(0)
                    clip_num+=1
                    new_img=np.concatenate((clip_img[None,...],clip_mask[None,...]),axis=0)
                    np.save(f"{save_path}/{name[:-4]}_{clip_num}.npy",new_img)

                    if(0):
                        cv2.imwrite("img.jpg",new_img[0])
                        cv2.imwrite("mask.jpg",new_img[1])
                        cv2.waitKey(0)
        print("----------------------------------")
